'''*
 * BILIBILI：技术宅物语
 * 河源创易电子有限公司
 *'''
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import torchvision
import numpy as np
import cv2
import socket
import time

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

COCO_INSTANCE_CATEGORY_NAMES = [
    '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
    'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
    'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A',
    'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
    'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table',
    'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
    'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book',
    'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]

def get_prediction(img_np, threshold):
    img = Image.fromarray(np.uint8(img_np))
    transform = T.Compose([T.ToTensor()]) 
    img = transform(img) 
    time_start = time.time()  # 记录开始时间
    pred = model([img]) 
    pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())] 
    pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]
    pred_score = list(pred[0]['scores'].detach().numpy())
    pred_t = -1;
    for x in pred_score:
        if x > threshold:
            pred_t = pred_score.index(x)
    pred_boxes = pred_boxes[:pred_t+1]
    pred_class = pred_class[:pred_t+1]
    print("pred_class:",pred_class)
    print("pred_boxes:",pred_boxes)
    return pred_boxes, pred_class

if __name__ == '__main__':
    usoc = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) #建一个UDP socket
    usoc.bind(('',10000)) #监听端口号
    print("Start!")
    frameIdNow = 0 #当前帧帧号
    frameSizeNow = 0 #当前接收到的帧体积
    packetIdNow = 0 #最新已接收的数据包号
    packetCount = 0 #当前帧包的数量
    packetLen = 0 #标准数据包大小
    frameSizeOk = 0 #当前帧已接收数据量
    jpgBuff = bytes('', 'utf-8') #图片数据缓存
    while True:
        udpbuff, address = usoc.recvfrom(10240) #阻塞接收数据，最多一次接收10240字节
        #解析数据
        frameId = (udpbuff[1] << 24) + (udpbuff[2] << 16) + (udpbuff[3] << 8) + udpbuff[4] #获取帧号
        frameSize = (udpbuff[5] << 24) + (udpbuff[6] << 16) + (udpbuff[7] << 8) + udpbuff[8] #获取帧体积
        packetId = udpbuff[10] #获取包号
        packetSize = (udpbuff[13] << 8) + udpbuff[14] #获取包体积

        if frameIdNow != frameId: #换帧，记录新一帧的数据信息
            frameIdNow = frameId	#更新帧号
            frameSizeNow = frameSize #更新帧体积
            packetCount = udpbuff[9] #更新数据包数量
            packetLen = (udpbuff[11] << 8) + udpbuff[12] #更新数据包长度
            frameSizeOk = 0  #清除当前帧已接收数据量
            packetIdNow = 0  #最新已接收数据包号清零
            jpgBuff = bytes('', 'utf-8') #清空图片数据缓存
        #复制至缓冲区，并只接收安全范围内的数据包
        if (packetId <= packetCount) and (packetId > packetIdNow): #新数据包包号不超过总数据包数量，且包号刚好比前一包多1
            if packetSize == (len(udpbuff)-15): #数据包减去包头等于包体积
                if (packetSize == packetLen) or (packetId == packetCount): #标准包或最后一包
                    jpgBuff = jpgBuff + udpbuff[15:] #拼接数据包
                    frameSizeOk = frameSizeOk + len(udpbuff) - 15 #帧数据总量累加
        if frameSizeNow == frameSizeOk: #当前帧接收完成
            nparr = np.frombuffer(jpgBuff, dtype=np.uint8) #将图片数组转为numpy数组
            image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) #解码图片
            
            time_start = time.clock()  # 记录开始时间
            boxes, pred_cls = get_prediction(image, 0.5) 
            time_end = time.clock()  # 记录结束时间
            time_sum = time_end - time_start  # 计算的时间差为程序的执行时间，单位为秒/s
            print(time_sum)

            rect_th=1
            text_size=1
            text_th=1
            for i in range(len(boxes)):
                cv2.rectangle(image, boxes[i][0], boxes[i][1],color=(0, 255, 0), thickness=rect_th)
                cv2.putText(image,pred_cls[i], boxes[i][0],  cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=text_th)
            cv2.imshow('ESP',image) #将图片显示出来
            if cv2.waitKey(1)==27: #按下ESC键退出
                break;
    usoc.close()
